{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.005475,
     "end_time": "2023-09-28T06:31:36.113528",
     "exception": false,
     "start_time": "2023-09-28T06:31:36.108053",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# CenterPose Model Training using Synthetic Data from Omniverse Replicator Objects Extension\n",
    "\n",
    "[NVIDIA Omniverse™ Isaac Sim](https://docs.omniverse.nvidia.com/isaacsim/latest/index.html) is a robotics simulation toolkit for the NVIDIA Omniverse™ platform. Isaac Sim has essential features for building virtual robotic worlds and experiments. It provides researchers and practitioners with the tools and workflows they need to create robust, physically accurate simulations and synthetic datasets.\n",
    "\n",
    "## What is Replicator Object Extension?\n",
    "\n",
    "[omni.replicator.object](https://docs.omniverse.nvidia.com/isaacsim/latest/replicator_tutorials/tutorial_replicator_object.html) is an extension that generates synthetic data for model training without requiring any changes to the code. It can be used for various tasks, such as retail object detection and robotics. The extension takes a YAML description file as input, which describes a mutable scene or a hierarchy of stacked description files. It then outputs a description file along with graphics content, including RGB images, 2D/3D bounding boxes, segmentation masks, and more.\n",
    "\n",
    "<img align=\"center\" src=\"https://docs.omniverse.nvidia.com/isaacsim/latest/_images/overview.png\" width=\"540\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.003275,
     "end_time": "2023-09-28T06:31:36.121470",
     "exception": false,
     "start_time": "2023-09-28T06:31:36.118195",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to generate a synthetic dataset for CenterPose TAO training using the Omniverse Replicator Objects extension. The objectives of this notebook are as follows:\n",
    "\n",
    "- Setting up a configuration file for generating a synthetic dataset in Omniverse.\n",
    "- Generating the synthetic training data based on the configuration file.\n",
    "- Running post-processing methods to adjust and verify the annotations.\n",
    "- Training the CenterPose model using TAO.\n",
    "\n",
    "By the end of this notebook, you will have generated a synthetic training set and a trained centerpose model.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.002853,
     "end_time": "2023-09-28T06:31:36.127828",
     "exception": false,
     "start_time": "2023-09-28T06:31:36.124975",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the generated synthetic dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/centerpose/results`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-28T06:31:36.134970Z",
     "iopub.status.busy": "2023-09-28T06:31:36.134429Z",
     "iopub.status.idle": "2023-09-28T06:31:36.151027Z",
     "shell.execute_reply": "2023-09-28T06:31:36.149850Z"
    },
    "papermill": {
     "duration": 0.022865,
     "end_time": "2023-09-28T06:31:36.153016",
     "exception": false,
     "start_time": "2023-09-28T06:31:36.130151",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tao-experiments\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\", \"centerpose\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"centerpose\", \"results\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/centerpose\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-28T06:31:36.158563Z",
     "iopub.status.busy": "2023-09-28T06:31:36.157979Z",
     "iopub.status.idle": "2023-09-28T06:31:37.143646Z",
     "shell.execute_reply": "2023-09-28T06:31:37.142270Z"
    },
    "papermill": {
     "duration": 0.992229,
     "end_time": "2023-09-28T06:31:37.147071",
     "exception": false,
     "start_time": "2023-09-28T06:31:36.154842",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-28T06:31:37.159684Z",
     "iopub.status.busy": "2023-09-28T06:31:37.159076Z",
     "iopub.status.idle": "2023-09-28T06:31:37.170468Z",
     "shell.execute_reply": "2023-09-28T06:31:37.169350Z"
    },
    "papermill": {
     "duration": 0.020526,
     "end_time": "2023-09-28T06:31:37.172894",
     "exception": false,
     "start_time": "2023-09-28T06:31:37.152368",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tao_configs = {\n",
    "   \"Mounts\":[\n",
    "         # Mapping the Local project directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       }\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         },\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid()),\n",
    "        \"network\": \"host\"\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tao_configs, mfile, indent=4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-28T06:31:37.181775Z",
     "iopub.status.busy": "2023-09-28T06:31:37.181330Z",
     "iopub.status.idle": "2023-09-28T06:31:37.516592Z",
     "shell.execute_reply": "2023-09-28T06:31:37.515009Z"
    },
    "papermill": {
     "duration": 0.343835,
     "end_time": "2023-09-28T06:31:37.520078",
     "exception": false,
     "start_time": "2023-09-28T06:31:37.176243",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.00481,
     "end_time": "2023-09-28T06:31:37.530385",
     "exception": false,
     "start_time": "2023-09-28T06:31:37.525575",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in the `nvidia-pyindex` python index. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python >=3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-28T06:31:37.539705Z",
     "iopub.status.busy": "2023-09-28T06:31:37.539066Z",
     "iopub.status.idle": "2023-09-28T06:31:40.612100Z",
     "shell.execute_reply": "2023-09-28T06:31:40.610503Z"
    },
    "papermill": {
     "duration": 3.081828,
     "end_time": "2023-09-28T06:31:40.615712",
     "exception": false,
     "start_time": "2023-09-28T06:31:37.533884",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-pyindex\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-09-28T06:31:40.632020Z",
     "iopub.status.busy": "2023-09-28T06:31:40.631670Z",
     "iopub.status.idle": "2023-09-28T06:31:41.324687Z",
     "shell.execute_reply": "2023-09-28T06:31:41.323092Z"
    },
    "papermill": {
     "duration": 0.701693,
     "end_time": "2023-09-28T06:31:41.328397",
     "exception": false,
     "start_time": "2023-09-28T06:31:40.626704",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.004888,
     "end_time": "2023-09-28T06:31:41.345632",
     "exception": false,
     "start_time": "2023-09-28T06:31:41.340744",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 2. Generate synthetic dataset and verify the annotations <a class=\"anchor\" id=\"head-2\"></a>\n",
    "### 2.1 Setup the configuration file for the Omniverse Replicator\n",
    "In this section, we will use the \"pallet\" dataset for the tutorial. The following script will automatically generate the synthetic \"pallet\" dataset:\n",
    "\n",
    "- Please note that all the USD contents are sourced from `omniverse://content.ov.nvidia.com/`.\n",
    "- Make sure you have the necessary access before running the code.\n",
    "- More details for the hyper-parameters can be found in the [omni.replicator.object](https://docs.omniverse.nvidia.com/isaacsim/latest/replicator_tutorials/tutorial_replicator_object.html) page.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file=\"\"\"\n",
    "omni.replicator.object:\n",
    "  version: 0.2.16\n",
    "  num_frames: 20\n",
    "  seed: 100\n",
    "  inter_frame_time: 1\n",
    "  gravity: 10000\n",
    "  position_H:\n",
    "    harmonizer_type: mutable_attribute\n",
    "    mutable_attribute:\n",
    "      distribution_type: range\n",
    "      start:\n",
    "      - -94.77713317047056\n",
    "      - 0\n",
    "      - -35.661244451558446\n",
    "      end:\n",
    "      - -94.77713317047056\n",
    "      - 0\n",
    "      - -35.661244451558446\n",
    "  screen_height: 720\n",
    "  focal_length: 14.228393962367306\n",
    "  output_path: /tmpsrc/results\n",
    "  horizontal_aperture: 20.955\n",
    "  screen_width: 1080\n",
    "  camera_parameters:\n",
    "    far_clip: 100000\n",
    "    focal_length: $(focal_length)\n",
    "    horizontal_aperture: $(horizontal_aperture)\n",
    "    near_clip: 0.1\n",
    "    screen_height: $(screen_height)\n",
    "    screen_width: $(screen_width)\n",
    "  default_camera:\n",
    "    count: 1\n",
    "    camera_parameters: $(camera_parameters)\n",
    "    transform_operators:\n",
    "    - translate_global:\n",
    "        distribution_type: harmonized\n",
    "        harmonizer_name: position_H\n",
    "    - rotateY: $[seed]*20\n",
    "    - rotateX:\n",
    "        distribution_type: range\n",
    "        start: -15\n",
    "        end: -25\n",
    "    - translate:\n",
    "        distribution_type: range\n",
    "        start:\n",
    "        - -40\n",
    "        - -30\n",
    "        - 400\n",
    "        end:\n",
    "        - 40\n",
    "        - 30\n",
    "        - 550\n",
    "    type: camera\n",
    "  distant_light:\n",
    "    color:\n",
    "      distribution_type: range\n",
    "      end:\n",
    "      - 1.3\n",
    "      - 1.3\n",
    "      - 1.3\n",
    "      start:\n",
    "      - 0.7\n",
    "      - 0.7\n",
    "      - 0.7\n",
    "    count: 5\n",
    "    intensity:\n",
    "      distribution_type: range\n",
    "      end: 600\n",
    "      start: 150\n",
    "    subtype: distant\n",
    "    transform_operators:\n",
    "    - rotateY:\n",
    "        distribution_type: range\n",
    "        end: 180\n",
    "        start: -180\n",
    "    - rotateX:\n",
    "        distribution_type: range\n",
    "        end: -10\n",
    "        start: -40\n",
    "    type: light\n",
    "  dome_light:\n",
    "    type: light\n",
    "    subtype: dome\n",
    "    color:\n",
    "      distribution_type: range\n",
    "      start:\n",
    "      - 0.7\n",
    "      - 0.7\n",
    "      - 0.7\n",
    "      end:\n",
    "      - 1.3\n",
    "      - 1.3\n",
    "      - 1.3\n",
    "    intensity:\n",
    "      distribution_type: range\n",
    "      start: 1000\n",
    "      end: 3000\n",
    "    transform_operators:\n",
    "    - rotateX: 270\n",
    "  plane:\n",
    "    physics: collision\n",
    "    type: geometry\n",
    "    subtype: plane\n",
    "    tracked: false\n",
    "    transform_operators:\n",
    "    - scale:\n",
    "      - 5\n",
    "      - 5\n",
    "      - 5\n",
    "  rotY_H:\n",
    "    harmonizer_type: mutable_attribute\n",
    "    mutable_attribute:\n",
    "      distribution_type: range\n",
    "      start: 0\n",
    "      end: 0\n",
    "  translate_H:\n",
    "    harmonizer_type: mutable_attribute\n",
    "    mutable_attribute:\n",
    "      distribution_type: range\n",
    "      start:\n",
    "      - 0\n",
    "      - 60\n",
    "      - 0\n",
    "      end:\n",
    "      - 0\n",
    "      - 30\n",
    "      - 0\n",
    "\n",
    "  pallet:\n",
    "    count: 2\n",
    "    physics: rigidbody\n",
    "    type: geometry\n",
    "    subtype: mesh\n",
    "    tracked: true\n",
    "    transform_operators:\n",
    "    - translate_global:\n",
    "        distribution_type: harmonized\n",
    "        harmonizer_name: position_H\n",
    "    - translate:\n",
    "      - 120 * ($[index]%2)\n",
    "      - 10 * ($[index]-1) * ($[index])\n",
    "      - 0\n",
    "    - rotateXYZ:\n",
    "      - -90\n",
    "      - 0\n",
    "      - 0\n",
    "    - scale:\n",
    "      - 1\n",
    "      - 1\n",
    "      - 1\n",
    "    usd_path:\n",
    "      distribution_type: set\n",
    "      values: \n",
    "      - omniverse://content.ov.nvidia.com/NVIDIA/Assets/DigitalTwin/Assets/Warehouse/Shipping/Pallets/Wood/Block_A/BlockPallet_A08_PR_NVD_01.usd\n",
    "  box:\n",
    "    count: 2\n",
    "    physics: rigidbody\n",
    "    type: geometry\n",
    "    subtype: mesh\n",
    "    tracked: false\n",
    "    transform_operators:\n",
    "    - translate_global:\n",
    "        distribution_type: harmonized\n",
    "        harmonizer_name: position_H\n",
    "    - translate_pallet:\n",
    "        distribution_type: harmonized\n",
    "        harmonizer_name: translate_H\n",
    "    - rotateY:\n",
    "        distribution_type: harmonized\n",
    "        harmonizer_name: rotY_H\n",
    "    - translate:\n",
    "      - 120 * ($[index])\n",
    "      - 20\n",
    "      - 0\n",
    "    - rotateXYZ:\n",
    "      - 0\n",
    "      - -90\n",
    "      - -90\n",
    "    - scale:\n",
    "      - 12\n",
    "      - 10\n",
    "      - 6\n",
    "    usd_path:\n",
    "      distribution_type: set\n",
    "      values:\n",
    "      - omniverse://content.ov.nvidia.com/NVIDIA/Assets/DigitalTwin/Assets/Warehouse/Shipping/Cardboard_Boxes/White_A/WhiteCorrugatedBox_A01_10x10x10cm_PR_NVD_01.usd\n",
    "      - omniverse://content.ov.nvidia.com/NVIDIA/Assets/DigitalTwin/Assets/Warehouse/Shipping/Cardboard_Boxes/Cube_A/CubeBox_A01_10cm_PR_NVD_01.usd\n",
    "  warehouse:\n",
    "    type: geometry\n",
    "    subtype: mesh\n",
    "    usd_path: omniverse://content.ov.nvidia.com/NVIDIA/Assets/Isaac/2023.1.1/Isaac/Environments/Simple_Warehouse/warehouse_with_forklifts.usd\n",
    "    transform_operators:\n",
    "    - translate:\n",
    "      - -200\n",
    "      - 0.1\n",
    "      - 0\n",
    "    - rotateXYZ:\n",
    "      - 0\n",
    "      - -90\n",
    "      - -90\n",
    "    - scale:\n",
    "      - 100\n",
    "      - 100\n",
    "      - 100\n",
    "\n",
    "  output_switches:\n",
    "    images: True\n",
    "    labels: True\n",
    "    descriptions: False\n",
    "    3d_labels: True\n",
    "    segmentation: False\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the configuration file and save to the local.\n",
    "import yaml\n",
    "yaml_file = yaml.safe_load(config_file)\n",
    "with open(os.path.join(os.getenv(\"HOST_DATA_DIR\", os.getcwd()), 'config.yaml'), 'w') as outfile:\n",
    "    yaml.dump(yaml_file, outfile, default_flow_style=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "912d080a",
   "metadata": {},
   "source": [
    "\n",
    "* Note that you could simple adjust the hyper-parameter from config file to generate the synthetic data on your use case.\n",
    "* **If you are using content from omniver nucleus servers you need to set your own OMNI_USER and OMNI_PASS variables.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the Omniverse password\n",
    "%env OMNI_PASS=YOUR_OWN_OMNI_PASSWORD"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37aa3fbb",
   "metadata": {},
   "source": [
    "### 2.2 Launch the synthetic data generation\n",
    "Launch the synthetic data generation with the Omniverse replicator object extension inside the container.\n",
    "\n",
    "The following example is one of the synthetic data generation scenes. You can modify the config file to generate different scenes with various objects, backgrounds, and numbers of target objects.\n",
    "* Note that the current synthetic data generation pipeline only supports single GPU. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8a2c56c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker run --gpus device=0 -it \\\n",
    "    --entrypoint /bin/bash \\\n",
    "    --network host \\\n",
    "    -v $HOST_DATA_DIR:/tmpsrc \\\n",
    "    -e OMNI_USER='$omni-api-token' \\\n",
    "    -e OMNI_PASS=$OMNI_PASS \\\n",
    "    nvcr.io/nvidia/isaac-sim:4.0.0 \\\n",
    "    -c \"apt-get update && apt-get install libglib2.0-dev -y && bash isaac-sim.sh --no-window --allow-root --/windowless=True --allow-root --/log/outputStreamLevel=fatal --/app/extensions/fastImporter/enabled=false --enable omni.replicator.object --/config/file=/tmpsrc/config.yaml\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31003ff5",
   "metadata": {},
   "source": [
    "### 2.3 Visualize the generated data\n",
    "In this section, we will use the `visualization` tool to visualize the generated synthetic data. The generation tool produces synthetic images and corresponding JSON files that contain training annotation information.\n",
    "\n",
    "Once the synthetic data is generated, it is stored in the `HOST_DATA_DIR`. It can be visualized using the following commands.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98cc25cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the matplotlib dependencies.\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\"\n",
    "!pip3 install opencv-python==4.8.0.74\n",
    "!pip3 install numpy==1.24.4\n",
    "!pip3 install pyrr\n",
    "!pip3 install scipy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec84014f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg']\n",
    "\n",
    "def visualize_images(output_path, num_cols=4, num_images=10):\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[40,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1995bbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "# Note that the sample spec is not meant to produce SOTA (state-of-the-art) accuracy on Objectron dataset.\n",
    "IMAGE_DIR = os.path.join(os.environ['HOST_DATA_DIR'], 'results', 'images')\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(IMAGE_DIR, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0713322f",
   "metadata": {},
   "source": [
    "### 2.4 Visualize and verify the synthetic data annotation \n",
    "In this section, we visualize the annotation for the generated synthetic data by using the related json file. \n",
    "- We first setup the PnP solver for verifying the 2d keypoints, camera intrinsic and the object scale. \n",
    "- We then visualize the annotation on the image. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "826879ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "from enum import IntEnum\n",
    "from pyrr import Quaternion\n",
    "from scipy.spatial.transform import Rotation as R\n",
    "\n",
    "class CuboidVertexType(IntEnum):\n",
    "    FrontTopRight = 0\n",
    "    FrontTopLeft = 1\n",
    "    FrontBottomLeft = 2\n",
    "    FrontBottomRight = 3\n",
    "    RearTopRight = 4\n",
    "    RearTopLeft = 5\n",
    "    RearBottomLeft = 6\n",
    "    RearBottomRight = 7\n",
    "    Center = 8\n",
    "    TotalCornerVertexCount = 8  # Corner vertexes doesn't include the center point\n",
    "    TotalVertexCount = 9\n",
    "\n",
    "class Cuboid3d():\n",
    "    '''This class contains a 3D cuboid.'''\n",
    "\n",
    "    # Create a box with a certain size\n",
    "    def __init__(self, size3d=[1.0, 1.0, 1.0],\n",
    "                 coord_system=None, parent_object=None):\n",
    "        # NOTE: This local coordinate system is similar\n",
    "        # to the intrinsic transform matrix of a 3d object\n",
    "        self.center_location = [0, 0, 0]\n",
    "        # self.center_location = [size3d[0]/2,size3d[1]/2,size3d[2]/2]\n",
    "        self.coord_system = coord_system\n",
    "        self.size3d = size3d\n",
    "        self._vertices = [0, 0, 0] * CuboidVertexType.TotalCornerVertexCount\n",
    "        # self._vertices = [0, 0, 0] * CuboidVertexType.TotalVertexCount\n",
    "\n",
    "        self.generate_vertexes()\n",
    "\n",
    "    def get_vertex(self, vertex_type):\n",
    "        \"\"\"Returns the location of a vertex.\n",
    "\n",
    "        Args:\n",
    "            vertex_type: enum of type CuboidVertexType\n",
    "\n",
    "        Returns:\n",
    "            Numpy array(3) - Location of the vertex type in the cuboid\n",
    "        \"\"\"\n",
    "        return self._vertices[vertex_type]\n",
    "\n",
    "    def get_vertices(self):\n",
    "        return self._vertices\n",
    "\n",
    "    def generate_vertexes(self):\n",
    "        width, height, depth = self.size3d\n",
    "\n",
    "        # By default just use the normal OpenCV coordinate system\n",
    "        if (self.coord_system is None):\n",
    "            cx, cy, cz = self.center_location\n",
    "            # X axis point to the right\n",
    "            right = cx + width / 2.0\n",
    "            left = cx - width / 2.0\n",
    "            # Y axis point upward\n",
    "            top = cy + height / 2.0\n",
    "            bottom = cy - height / 2.0\n",
    "            # Z axis point forward\n",
    "            front = cz + depth / 2.0\n",
    "            rear = cz - depth / 2.0\n",
    "\n",
    "            # List of 8 vertices of the box\n",
    "            self._vertices = [\n",
    "                # self.center_location,   # Center\n",
    "                [left, bottom, rear],  # Rear Bottom Left\n",
    "                [left, bottom, front],  # Front Bottom Left\n",
    "                [left, top, rear],  # Rear Top Left\n",
    "                [left, top, front],  # Front Top Left\n",
    "\n",
    "                [right, bottom, rear],  # Rear Bottom Right\n",
    "                [right, bottom, front],  # Front Bottom Right\n",
    "                [right, top, rear],  # Rear Top Right\n",
    "                [right, top, front],  # Front Top Right\n",
    "\n",
    "            ]\n",
    "\n",
    "class CuboidPNPSolver(object):\n",
    "    \"\"\"\n",
    "    This class is used to find the 6-DoF pose of a cuboid given its projected vertices.\n",
    "\n",
    "    Runs perspective-n-point (PNP) algorithm.\n",
    "    \"\"\"\n",
    "\n",
    "    # Class variables\n",
    "    cv2version = cv2.__version__.split('.')\n",
    "    cv2majorversion = int(cv2version[0])\n",
    "\n",
    "    def __init__(self, scaling_factor=1,\n",
    "                 camera_intrinsic_matrix=None,\n",
    "                 cuboid3d=None,\n",
    "                 dist_coeffs=np.zeros((4, 1)),\n",
    "                 min_required_points=4\n",
    "                 ):\n",
    "\n",
    "        self.min_required_points = max(4, min_required_points)\n",
    "        self.scaling_factor = scaling_factor\n",
    "\n",
    "        if (not camera_intrinsic_matrix is None):\n",
    "            self._camera_intrinsic_matrix = camera_intrinsic_matrix\n",
    "        else:\n",
    "            self._camera_intrinsic_matrix = np.array([\n",
    "                [0, 0, 0],\n",
    "                [0, 0, 0],\n",
    "                [0, 0, 0]\n",
    "            ])\n",
    "        self._cuboid3d = cuboid3d\n",
    "        \n",
    "        self._dist_coeffs = dist_coeffs\n",
    "        \n",
    "    def set_camera_intrinsic_matrix(self, new_intrinsic_matrix):\n",
    "        '''Sets the camera intrinsic matrix'''\n",
    "        self._camera_intrinsic_matrix = new_intrinsic_matrix\n",
    "\n",
    "    def set_dist_coeffs(self, dist_coeffs):\n",
    "        '''Sets the camera intrinsic matrix'''\n",
    "        self._dist_coeffs = dist_coeffs\n",
    "\n",
    "    def solve_pnp(self,\n",
    "                  cuboid2d_points,\n",
    "                  pnp_algorithm=None,\n",
    "                  OPENCV_RETURN = True,\n",
    "                  verbose = False\n",
    "                  ):\n",
    "        \"\"\"\n",
    "        Detects the rotation and traslation \n",
    "        of a cuboid object from its vertexes' \n",
    "        2D location in the image\n",
    "\n",
    "        Inputs:\n",
    "        - cuboid2d_points:  list of XY tuples\n",
    "        - pnp_algorithm: None  \n",
    "          ...\n",
    "\n",
    "        Outputs:\n",
    "        - location in 3D\n",
    "        - pose in 3D (as quaternion)\n",
    "        - projected points:  np.ndarray of np.ndarrays\n",
    "\n",
    "        \"\"\"\n",
    "        \n",
    "        # Fallback to default PNP algorithm base on OpenCV version\n",
    "        if pnp_algorithm is None:\n",
    "            if CuboidPNPSolver.cv2majorversion == 2:\n",
    "                pnp_algorithm = cv2.CV_ITERATIVE\n",
    "            elif CuboidPNPSolver.cv2majorversion == 3:\n",
    "                pnp_algorithm = cv2.SOLVEPNP_ITERATIVE\n",
    "            else:\n",
    "                pnp_algorithm = cv2.SOLVEPNP_ITERATIVE\n",
    "\n",
    "        location = None\n",
    "        quaternion = None\n",
    "        location_new = None\n",
    "        quaternion_new = None\n",
    "        reprojectionError = None\n",
    "\n",
    "        projected_points = cuboid2d_points\n",
    "        cuboid3d_points = np.array(self._cuboid3d.get_vertices())\n",
    "        \n",
    "        obj_2d_points = []\n",
    "        obj_3d_points = []\n",
    "\n",
    "        # 8*n points\n",
    "        for i in range(len(cuboid2d_points)):\n",
    "            check_point_2d = cuboid2d_points[i]\n",
    "            # Ignore invalid points\n",
    "            if (check_point_2d is None or check_point_2d[0] < -5000 or check_point_2d[1] < -5000):\n",
    "                continue\n",
    "            obj_2d_points.append(check_point_2d)\n",
    "            obj_3d_points.append(\n",
    "                cuboid3d_points[int(i // (len(cuboid2d_points) / CuboidVertexType.TotalCornerVertexCount))]) # TotalCornerVertexCount = 8 \n",
    "        \n",
    "        obj_2d_points = np.array(obj_2d_points, dtype=float)\n",
    "        obj_3d_points = np.array(obj_3d_points, dtype=float)\n",
    "        valid_point_count = len(obj_2d_points)\n",
    "\n",
    "        # Can only do PNP if we have more than 3 valid points\n",
    "        is_points_valid = valid_point_count >= self.min_required_points\n",
    "\n",
    "        if is_points_valid:\n",
    "\n",
    "            # Heatmap representation may have less than 6 points, in which case we have to use another pnp algorithm\n",
    "            if valid_point_count < 6:\n",
    "                pnp_algorithm = cv2.SOLVEPNP_EPNP\n",
    "\n",
    "            # Usually, we use this one\n",
    "            ret, rvec, tvec, reprojectionError = cv2.solvePnPGeneric(\n",
    "                obj_3d_points,\n",
    "                obj_2d_points,\n",
    "                self._camera_intrinsic_matrix,\n",
    "                self._dist_coeffs,\n",
    "                flags=pnp_algorithm\n",
    "            )\n",
    "            \n",
    "            if ret:\n",
    "\n",
    "                rvec = np.array(rvec[0])\n",
    "                tvec = np.array(tvec[0])\n",
    "\n",
    "                reprojectionError = reprojectionError.flatten()[0]\n",
    "\n",
    "                # Convert OpenCV coordinate system to OpenGL coordinate system\n",
    "                transformation = np.identity(4)\n",
    "                r = R.from_rotvec(rvec.reshape(1, 3))\n",
    "                transformation[:3, :3] = r.as_matrix()\n",
    "                transformation[:3, 3] = tvec.reshape(1, 3)\n",
    "                M = np.zeros((4, 4))\n",
    "                M[0, 1] = 1\n",
    "                M[1, 0] = 1\n",
    "                M[3, 3] = 1\n",
    "                M[2, 2] = -1\n",
    "                transformation = np.matmul(M, transformation)\n",
    "\n",
    "                rvec_new = R.from_matrix(transformation[:3, :3]).as_rotvec()\n",
    "                tvec_new = transformation[:3, 3]\n",
    "\n",
    "                # OpenGL result, to be compared against GT\n",
    "                location_new = list(x for x in tvec_new)\n",
    "                quaternion_new = self.convert_rvec_to_quaternion(rvec_new)\n",
    "\n",
    "                # OpenCV result\n",
    "                location = list(x[0] for x in tvec)\n",
    "                quaternion = self.convert_rvec_to_quaternion(rvec)\n",
    "                \n",
    "                # Still use OpenCV way to project 3D points\n",
    "                projected_points, _ = cv2.projectPoints(cuboid3d_points, rvec, tvec, self._camera_intrinsic_matrix,\n",
    "                                                        self._dist_coeffs)\n",
    "                \n",
    "                projected_points = np.squeeze(projected_points)\n",
    "                \n",
    "                x, y, z = location\n",
    "                if z < 0:\n",
    "                    location = None\n",
    "                    quaternion = None\n",
    "                    location_new = None\n",
    "                    quaternion_new = None\n",
    "\n",
    "                    if verbose:\n",
    "                        print(\"PNP solution is behind the camera (Z < 0) => Fail\")\n",
    "                else:\n",
    "                    if verbose:\n",
    "                        print(\"solvePNP found good results - location: {} - rotation: {} !!!\".format(location, quaternion))\n",
    "            else:\n",
    "                if verbose:\n",
    "                    print('Error:  solvePnP return false ****************************************')\n",
    "        else:\n",
    "            if verbose:\n",
    "                print(\"Need at least 4 valid points in order to run PNP. Currently: {}\".format(valid_point_count))\n",
    "\n",
    "        if OPENCV_RETURN:\n",
    "            # Return OpenCV result for demo\n",
    "            return location, quaternion, projected_points, reprojectionError\n",
    "        else:\n",
    "            # Return OpenGL result for eval\n",
    "            return location_new, quaternion_new, projected_points, reprojectionError\n",
    "        \n",
    "    def convert_rvec_to_quaternion(self, rvec): \n",
    "        '''Convert rvec (which is log quaternion) to quaternion'''\n",
    "        theta = np.sqrt(rvec[0] * rvec[0] + rvec[1] * rvec[1] + rvec[2] * rvec[2])  # in radians\n",
    "        raxis = [rvec[0] / theta, rvec[1] / theta, rvec[2] / theta]\n",
    "        return Quaternion.from_axis_rotation(raxis, theta)\n",
    "\n",
    "def pnp_processing(points_filtered, scale, cam_intrinsic, OPENCV_RETURN=True):\n",
    "    # initial a 3d cuboid\n",
    "    cuboid3d = Cuboid3d(1 * np.array(scale))\n",
    "\n",
    "    pnp_solver = \\\n",
    "        CuboidPNPSolver(\n",
    "            cuboid3d=cuboid3d\n",
    "        )\n",
    "    pnp_solver.set_camera_intrinsic_matrix(cam_intrinsic)\n",
    "    \n",
    "    location, quaternion, projected_points, reprojectionError = pnp_solver.solve_pnp(\n",
    "        points_filtered, OPENCV_RETURN=OPENCV_RETURN)  # N * 2\n",
    "    \n",
    "    if location is not None:\n",
    "\n",
    "        ori = R.from_quat(quaternion).as_matrix()\n",
    "        pose_pred = np.identity(4)\n",
    "        pose_pred[:3, :3] = ori\n",
    "        pose_pred[:3, 3] = location\n",
    "        point_3d_obj = cuboid3d.get_vertices()\n",
    "        \n",
    "        point_3d_cam = pose_pred @ np.hstack(\n",
    "            (np.array(point_3d_obj), np.ones((np.array(point_3d_obj).shape[0], 1)))).T\n",
    "        point_3d_cam = point_3d_cam[:3, :].T  # 8 * 3\n",
    "        \n",
    "        \n",
    "        # Add the centroid\n",
    "        point_3d_cam = np.insert(point_3d_cam, 0, np.mean(point_3d_cam, axis=0), axis=0)\n",
    "        \n",
    "        # Add the center\n",
    "        projected_points = np.insert(projected_points, 0, np.mean(projected_points, axis=0), axis=0)\n",
    "        \n",
    "        return projected_points, point_3d_cam, location, quaternion\n",
    "\n",
    "    return None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff9f6c12",
   "metadata": {},
   "source": [
    "- Note that the default image format is `.jpg` from the replicator. Please revise the following image extension if you generated other image format. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6677463d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import cv2\n",
    "valid_image_ext = ['.jpg']\n",
    "\n",
    "def add_obj_order(img, keypoints2d, img_id='default', pred_flag='pred'):\n",
    "    bbox = np.array(keypoints2d, dtype=np.int32)\n",
    "    font = cv2.FONT_HERSHEY_SIMPLEX\n",
    "\n",
    "    for i in range(len(bbox)):\n",
    "        txt = '{:d}'.format(i)\n",
    "        cat_size = cv2.getTextSize(txt, font, 1, 2)[0]\n",
    "        cv2.putText(img, txt, (bbox[i][0], bbox[i][1] + cat_size[1]),\n",
    "                    font, 1, (0, 0, 255), thickness=2, lineType=cv2.LINE_AA)\n",
    "\n",
    "def add_coco_hp(img, points):\n",
    "    # use for drawing the 3D bounding box\n",
    "    edges = [[2, 4], [2, 6], [6, 8], [4, 8],\n",
    "                [1, 2], [3, 4], [5, 6], [7, 8],\n",
    "                [1, 3], [1, 5], [3, 7], [5, 7]]\n",
    "\n",
    "    num_joints = 8\n",
    "    points = np.array(points, dtype=np.int32).reshape(num_joints, 2)\n",
    "    # Draw edges\n",
    "    for j, e in enumerate(edges):\n",
    "        temp = [e[0] - 1, e[1] - 1]\n",
    "        edge_color = (0, 255, 0)  # bgr\n",
    "        if points[temp[1], 0] <= -10000 or points[temp[1], 1] <= -10000 or points[temp[0], 0] <= -10000 or \\\n",
    "                points[temp[0], 1] <= -10000:\n",
    "            continue\n",
    "        else:\n",
    "            cv2.line(img, (points[temp[0], 0], points[temp[0], 1]),\n",
    "                        (points[temp[1], 0], points[temp[1], 1]), edge_color, 2)\n",
    "\n",
    "def add_axes(img, box, cam_intrinsic):\n",
    "    # box 9x3 array\n",
    "    # OpenCV way\n",
    "    N = 20\n",
    "    # Centroid, top, front, right\n",
    "    axes_point_list = [0, box[3] - box[1], box[2] - box[1], box[5] - box[1]]\n",
    "    \n",
    "    viewport_point_list = []\n",
    "    for axes_point in axes_point_list:\n",
    "        vector = axes_point\n",
    "        vector = vector / np.linalg.norm(vector) * N if np.linalg.norm(vector) != 0 else 0\n",
    "        vector = vector + box[0]\n",
    "        vector = vector.flatten()\n",
    "\n",
    "        k_3d = np.array([vector[0], vector[1], vector[2]])\n",
    "        pp = np.matmul(cam_intrinsic, k_3d.reshape(3, 1))\n",
    "        viewport_point = [pp[0] / pp[2], pp[1] / pp[2]]\n",
    "        viewport_point_list.append((int(viewport_point[0]), int(viewport_point[1])))\n",
    "\n",
    "    # BGR space\n",
    "    cv2.line(img, viewport_point_list[0], viewport_point_list[1], (0, 255, 0), 5)  # y-> green\n",
    "    cv2.line(img, viewport_point_list[0], viewport_point_list[2], (255, 0, 0), 5)  # z-> blue\n",
    "    cv2.line(img, viewport_point_list[0], viewport_point_list[3], (0, 0, 255), 5)  # x-> red\n",
    "\n",
    "def visualize_annotation(imgs, anns, num_cols=4, num_images=10):\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[40, 30])\n",
    "    f.tight_layout()\n",
    "    for idx, img_path in enumerate(imgs[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "\n",
    "        with open(anns[idx]) as f:\n",
    "            ann = json.load(f)\n",
    "\n",
    "        num_objs = len(ann['objects'])\n",
    "\n",
    "        cx = ann['camera_data']['intrinsics']['cx']\n",
    "        cy = ann['camera_data']['intrinsics']['cy']\n",
    "        fx = ann['camera_data']['intrinsics']['fx']\n",
    "        fy = ann['camera_data']['intrinsics']['fy']\n",
    "        cam_intrinsic = np.array(\n",
    "                    [[fx, 0, cx], [0, fy, cy], [0, 0, 1]])\n",
    "\n",
    "        for k in range(num_objs):\n",
    "\n",
    "            # Read image and annotation\n",
    "            ann_obj = ann['objects'][k]\n",
    "\n",
    "            pts_ori = np.array(ann_obj['projected_cuboid'])\n",
    "            scale = np.array(ann_obj['scale'])\n",
    "            \n",
    "            # Run the PnP solver for futher verifying the 2d keypoint, scale and camera intrinsic. \n",
    "            projected_points, point_3d_cam, rotation, quaternion = pnp_processing(pts_ori[1:], scale, cam_intrinsic)\n",
    "\n",
    "            # GT visualization\n",
    "            add_obj_order(img, projected_points)\n",
    "            add_coco_hp(img, projected_points[1:])\n",
    "            add_axes(img, point_3d_cam, cam_intrinsic)\n",
    "\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16d345f5",
   "metadata": {},
   "source": [
    "The following visualization code is running the PnP solver to verify the 2d keypoint, scale and camera intrinsic with the pose. \n",
    "-  \"2D Keypoints\" + \"Object scale\" + \"Camera Intrinsic\" ==> PnP solver ==> Object Pose + projected keypoint + 3D keypoint in camera space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ef949ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "# Note that the sample spec is not meant to produce SOTA (state-of-the-art) accuracy on Objectron dataset.\n",
    "IMAGE_DIR = os.path.join(os.environ['HOST_DATA_DIR'], 'results', 'images')\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "imgs = [os.path.join(IMAGE_DIR, image) for image in os.listdir(IMAGE_DIR) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "anns = [image.replace(os.path.splitext(image)[1].lower(), '.json').replace('/images/', '/3d_labels/')\n",
    "        for image in imgs]\n",
    "visualize_annotation(imgs, anns, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86f4f80d",
   "metadata": {},
   "source": [
    "## 3. Post-processing Procedure before Launching the CenterPose Training\n",
    "In this section, we run some dataset post-processing for finalizing the annotations. \n",
    "- Note that this part will include into the pose writer in the next release.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08216acb",
   "metadata": {},
   "source": [
    "### 3.1 Rotate the annotation axis and add the missing hyperparameters\n",
    "\n",
    "In this section, we convert the annotation file to the [Objectron format](https://github.com/google-research-datasets/Objectron/blob/master/notebooks/Parse%20Annotations.ipynb). \n",
    "\n",
    "The Objectron format defines the orientation of the coordinate system as follows: the positive y-axis points upwards (aligned with gravity), the positive z-axis points towards the user, and the positive x-axis follows the right-hand rule. The front face is defined as the positive z-axis on the xy-plane, and the top face is defined as the positive y-axis on the xz-plane. Therefore, we need to convert the annotation to match the Objectron format.\n",
    "\n",
    "The original 3D label annotation is missing the image width, image height, and the camera view matrix. Therefore, we added these hyperparameters to the annotation file as well.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f15ff7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "Z_TO_Y_AXIS = [0, 2, 4, 1, 3, 6, 8, 5, 7]\n",
    "SCALE_CONV = [0, 2, 1]\n",
    "\n",
    "def annotation_converter(output_path):\n",
    "    imgs = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    anns = [image.replace(os.path.splitext(image)[1].lower(), '.json').replace('/images/', '/3d_labels/')\n",
    "            for image in imgs]\n",
    "\n",
    "    for idx, img_path in enumerate(imgs):\n",
    "\n",
    "        with open(anns[idx]) as f:\n",
    "            ann = json.load(f)\n",
    "\n",
    "        img = cv2.imread(img_path)\n",
    "        height, width, _ = img.shape\n",
    "        if \"width\" not in ann[\"camera_data\"]:\n",
    "            ann[\"camera_data\"][\"width\"] = width\n",
    "        if \"height\" not in ann[\"camera_data\"]:\n",
    "            ann[\"camera_data\"][\"height\"] = height\n",
    "        ann[\"camera_data\"][\"camera_view_matrix\"] = np.eye(4).tolist()\n",
    "        num_objs = len(ann['objects'])\n",
    "\n",
    "        dict_out = {\"camera_data\": ann[\"camera_data\"], \"objects\": []}\n",
    "        for k in range(num_objs):\n",
    "\n",
    "            # Read image and annotation\n",
    "            ann_obj = ann['objects'][k]\n",
    "\n",
    "            pts_ori = np.array(ann_obj['projected_cuboid'])\n",
    "            scale = np.array(ann_obj['scale'])\n",
    "            pts_3d = np.array(ann_obj['keypoints_3d'])\n",
    "            \n",
    "            # Rotate the 3D and 2D axis\n",
    "            new_keypoints_2d = []\n",
    "            new_keypoints_3d = []\n",
    "            new_keypoints_3d = pts_3d[Z_TO_Y_AXIS][Z_TO_Y_AXIS][Z_TO_Y_AXIS]\n",
    "            new_keypoints_2d = pts_ori[Z_TO_Y_AXIS][Z_TO_Y_AXIS][Z_TO_Y_AXIS]\n",
    "            new_scale = scale[SCALE_CONV]\n",
    "\n",
    "            ann_obj['projected_cuboid'] = new_keypoints_2d.tolist()\n",
    "            ann_obj['keypoints_3d'] = new_keypoints_3d.tolist()\n",
    "            ann_obj['scale'] = list(new_scale)\n",
    "            ann_obj['visibility'] = 1.0\n",
    "            dict_out[\"objects\"].append(ann_obj)\n",
    "\n",
    "        json_save = img_path.replace(os.path.splitext(img_path)[1], '.json')\n",
    "        with open(json_save, 'w+') as fp:\n",
    "            print(f\"Saving the updated 3d labels to {json_save}\")\n",
    "            json.dump(dict_out, fp, indent=4, sort_keys=True)\n",
    "\n",
    "annotation_converter(IMAGE_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91a4280e",
   "metadata": {},
   "source": [
    "### 3.2 AR data calcuation\n",
    "The CenterPose predicts the 3D bounding boxes from RGB images, the 3D keypoints may up to scale. However, the ground truth is at metric scale. In other words, they are may in different scale, because predicting scale from monocular camera (single RGB image) is ambiguous. Therefore, there has a method to re-scale the predicted box using ground planes (assuming the box is sitting on the ground)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "868d4608",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the Box function to calculate the ground planes.\n",
    "FACES = np.array([\n",
    "    [5, 6, 8, 7],  # +x on yz plane\n",
    "    [1, 3, 4, 2],  # -x on yz plane\n",
    "    [3, 7, 8, 4],  # +y on xz plane = top\n",
    "    [1, 2, 6, 5],  # -y on xz plane\n",
    "    [2, 4, 8, 6],  # +z on xy plane = front\n",
    "    [1, 5, 7, 3],  # -z on xy plane\n",
    "])\n",
    "\n",
    "class Box(object):\n",
    "  \"\"\"General 3D Oriented Bounding Box.\"\"\"\n",
    "\n",
    "  def __init__(self, vertices=None):\n",
    "\n",
    "    self.vertices = vertices\n",
    "\n",
    "  def get_ground_plane(self, gravity_axis=1):\n",
    "    \"\"\"Get ground plane under the box.\"\"\"\n",
    "\n",
    "    gravity = np.zeros(3)\n",
    "    gravity[gravity_axis] = 1\n",
    "\n",
    "    def get_face_normal(face, center):\n",
    "      \"\"\"Get a normal vector to the given face of the box.\"\"\"\n",
    "      v1 = self.vertices[face[0], :] - center\n",
    "      v2 = self.vertices[face[1], :] - center\n",
    "      normal = np.cross(v1, v2)\n",
    "      return normal\n",
    "\n",
    "    def get_face_center(face):\n",
    "      \"\"\"Get the center point of the face of the box.\"\"\"\n",
    "      center = np.zeros(3)\n",
    "      for vertex in face:\n",
    "        center += self.vertices[vertex, :]\n",
    "      center /= len(face)\n",
    "      return center\n",
    "\n",
    "    ground_plane_id = 0\n",
    "    ground_plane_error = 10.\n",
    "\n",
    "    # The ground plane is defined as a plane aligned with gravity.\n",
    "    # gravity is the (0, 1, 0) vector in the world coordinate system.\n",
    "    for i in [0, 2, 4]:\n",
    "      face = FACES[i, :]\n",
    "      center = get_face_center(face)\n",
    "      normal = get_face_normal(face, center)\n",
    "      w = np.cross(gravity, normal)\n",
    "      w_sq_norm = np.linalg.norm(w)\n",
    "      if w_sq_norm < ground_plane_error:\n",
    "        ground_plane_error = w_sq_norm\n",
    "        ground_plane_id = i\n",
    "\n",
    "    face = FACES[ground_plane_id, :]\n",
    "    center = get_face_center(face)\n",
    "    normal = get_face_normal(face, center)\n",
    "\n",
    "    # For each face, we also have a parallel face that it's normal is also\n",
    "    # aligned with gravity vector. We pick the face with lower height (y-value).\n",
    "    # The parallel to face 0 is 1, face 2 is 3, and face 4 is 5.\n",
    "    parallel_face_id = ground_plane_id + 1\n",
    "    parallel_face = FACES[parallel_face_id]\n",
    "    parallel_face_center = get_face_center(parallel_face)\n",
    "    parallel_face_normal = get_face_normal(parallel_face, parallel_face_center)\n",
    "    if parallel_face_center[gravity_axis] < center[gravity_axis]:\n",
    "      center = parallel_face_center\n",
    "      normal = parallel_face_normal\n",
    "    return center, normal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beb78cf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ar_data_converter(output_path):\n",
    "    imgs = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    anns = [image.replace(os.path.splitext(image)[1].lower(), '.json')\n",
    "            for image in imgs]\n",
    "\n",
    "    for idx, img_path in enumerate(imgs):\n",
    "\n",
    "        with open(anns[idx]) as f:\n",
    "            ann = json.load(f)\n",
    "\n",
    "        num_objs = len(ann['objects'])\n",
    "\n",
    "        dict_out = {\"camera_data\": ann[\"camera_data\"], \"objects\": ann[\"objects\"], \"AR_data\":{}}\n",
    "\n",
    "        for k in range(num_objs):\n",
    "\n",
    "            # Read image and annotation\n",
    "            ann_obj = ann['objects'][k]\n",
    "            ann_obj['AR_data'] = {}\n",
    "            pts_3d = np.array(ann_obj['keypoints_3d'])\n",
    "            \n",
    "            # Calculate the AR data\n",
    "            pts3d_box = Box(pts_3d)\n",
    "            center, normal = pts3d_box.get_ground_plane()\n",
    "\n",
    "            # setup the object-wise ar data values\n",
    "            ann_obj['AR_data']['plane_center'] = center.tolist()\n",
    "            ann_obj['AR_data']['plane_normal'] = normal.tolist()\n",
    "\n",
    "            # setup the global ar data values\n",
    "            dict_out['AR_data']['plane_center'] = center.tolist()\n",
    "            dict_out['AR_data']['plane_normal'] = normal.tolist()\n",
    "\n",
    "        json_save = img_path.replace(os.path.splitext(img_path)[1], '.json')\n",
    "        with open(json_save, 'w+') as fp:\n",
    "            print(f\"Saving the updated 3d labels to {json_save}\")\n",
    "            json.dump(dict_out, fp, indent=4, sort_keys=True)\n",
    "\n",
    "ar_data_converter(IMAGE_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fea0201c",
   "metadata": {},
   "source": [
    "## 4. Launch the CenterPose TAO Training using the generated synthetic dataset\n",
    "In this section, it will introduce how to use the synthetic dataset to launch the training. \n",
    "More details regarding to the hyper-parameters and the end-to-end training pipeline could be found in the [CenterPose Notebook](https://github.com/NVIDIA/tao_tutorials/blob/main/notebooks/tao_launcher_starter_kit/centerpose/centerpose.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Download the pre-trained model\n",
    "We will use NGC CLI to get the pre-trained models. For more details, go to [ngc.nvidia.com](ngc.nvidia.com) and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Installing NGC CLI on the local machine.\n",
    "# ## Download and install\n",
    "# %env CLI=ngccli_cat_linux.zip\n",
    "# !mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# # Remove any previously existing CLI installations\n",
    "# !rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "# !wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "# !unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "# !rm $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "# os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!mkdir -p $HOST_RESULTS_DIR/pretrained_models\n",
    "!ngc registry model download-version \"nvidia/tao/centerpose:trainable_fan_small\" --dest $HOST_RESULTS_DIR/pretrained_models\n",
    "print(\"Check if model is downloaded into dir.\")\n",
    "!ls -l $HOST_RESULTS_DIR/pretrained_models/centerpose_vtrainable_fan_small/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    "\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "    * train_data: images and annotation files for train data. Required to have correct camera calibration data\n",
    "    * val_data: images and annotation files for validation data. Required to have correct camera calibration data\n",
    "    * num_classes: number of categories, default is 1. The CenterPose is a category-based method\n",
    "    * batch_size: batch size for dataloader\n",
    "    * workers: number of workers to do data loading\n",
    "    * category: category name of the training object\n",
    "    * num_symmetry: number of symmetric rotations for the specific categories, e.g. bottle\n",
    "    * max_objs: maximum number of training objects in one image\n",
    "* model: configure the model setting\n",
    "    * down_ratio: down sample ratio for the input image, default is 4\n",
    "    * use_pretrained: flag to enable using the pretrained weights\n",
    "    * model_type: backbone types of the CenterPose, including FAN-variants and the DLA34 backbone\n",
    "    * pretrained_backbone_path: path to the pretrained backbone model. FAN-variants is supported. DLA34 backbone loads the pretrained weight automatically. \n",
    "* train: configure the training hyperparameters\n",
    "    * num_gpus: number of gpus \n",
    "    * validation_interval: validation interval\n",
    "    * checkpoint_interval: interval of saving the checkpoint\n",
    "    * num_epochs: number of epochs\n",
    "    * clip_grad_val: the value of cliping the gradient, default is 100.0\n",
    "    * randomseed: random seed for reproducing the accuracy\n",
    "    * resume_training_checkpoint_path: resume the training from the checkpoint path\n",
    "    * precision: If set to fp16, the training is run on Automatic Mixed Precision (AMP)\n",
    "    * optim:\n",
    "        * lr: learning rate for training the model\n",
    "        * lr_steps: learning rate decay step milestone (MultiStep)\n",
    "\n",
    "Please refer to the TAO documentation about CenterPose to get all the parameters that are configurable.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "!cat $HOST_SPECS_DIR/train.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.3 Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* Evaluation mainly uses 3D IoU and 2D MPE (mean pixel errors) metrics. For more info, please refer to: https://github.com/google-research-datasets/Objectron\n",
    "* For this demonstration, we set the training epoch equals to 1 so that the training can be completed faster.\n",
    "* Unlike the [original CenterPose paper](https://arxiv.org/abs/2109.06161), we also provided a more advanced backbone called [FAN](https://arxiv.org/abs/2204.12451) that has proven to achieve higher downstream results compared to DLA34. \n",
    "* If you wish to speed up training, you may try to set `train.precision=fp16` for mixed precision training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env MODEL_DIR = /model\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!echo $HOST_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"For multi-GPU, change train.num_gpus in train.yaml based on your machine.\")\n",
    "# If you face out of memory issue, you may reduce the batch size in the spec file by passing dataset. batch_size=2\n",
    "!tao model centerpose train \\\n",
    "          -e $SPECS_DIR/train_synthetic.yaml \\\n",
    "          results_dir=$RESULTS_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Trained checkpoints:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"train/centerpose_model_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/train/centerpose_model.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train/centerpose_model.pth"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48e82698",
   "metadata": {},
   "source": [
    "\n",
    "This notebook has come to an end.\n",
    "More details of end-to-end training and inference pipeline could be found in the [CenterPose Notebook](https://github.com/NVIDIA/tao_tutorials/blob/main/notebooks/tao_launcher_starter_kit/centerpose/centerpose.ipynb)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 6.348297,
   "end_time": "2023-09-28T06:31:41.570413",
   "environment_variables": {},
   "exception": null,
   "input_path": "../../cv/resource/notebooks/tao_launcher_starter_kit/centerpose/centerpose.ipynb",
   "output_path": "../../cv/resource/notebooks/tao_launcher_starter_kit/centerpose/centerpose.ipynb",
   "parameters": {},
   "start_time": "2023-09-28T06:31:35.222116",
   "version": "2.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
